Trapping rain water II

Time: O(MxNx(LogM+LogN)); Space: O(MxN); hard

Given an m x n matrix of positive integers representing the height of each unit cell in a 2D elevation map,

compute the volume of water it is able to trap after raining.

Example 1:

Input: heightMap =

[
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]

Output: 4

Explanation:

  • Before the rain:

After the rain, water is trapped between the blocks. The total volume of water trapped is 4.

Example 2:

Input: heightMap =

[
    [12,13,0,12],
    [13,4,13,12],
    [13,8,10,12],
    [12,13,12,12],
    [13,13,13,13]
]

Output: 14

Example 3:

Input: heightMap =

[
    [2,2,2,2],
    [2,2,3,4],
    [3,3,3,1],
    [2,3,4,5]
]

Output: 0

Constraints:

  • 1 <= m, n <= 110

  • 0 <= heightMap[i][j] <= 20000

[3]:

from heapq import heappush, heappop
class Solution1(object):
    """
    Time: O(M*N*Log(M+N))~O(M*N*Log(M*N))
    Space: O(M*N)
    """
    def trapRainWater(self, heightMap):
        """
        :type heightMap: List[List[int]]
        :rtype: int
        """
        m = len(heightMap)
        if not m:
            return 0
        n = len(heightMap[0])
        if not n:
            return 0

        is_visited = [[False for i in range(n)] for j in range(m)]

        heap = []

        for i in range(m):
            heappush(heap, [heightMap[i][0], i, 0])
            is_visited[i][0] = True
            heappush(heap, [heightMap[i][n-1], i, n-1])
            is_visited[i][n-1] = True

        for j in range(n):
            heappush(heap, [heightMap[0][j], 0, j])
            is_visited[0][j] = True
            heappush(heap, [heightMap[m-1][j], m-1, j])
            is_visited[m-1][j] = True

        trap = 0
        while heap:
            height, i, j = heappop(heap)
            for (dx, dy) in [(1,0), (-1,0), (0,1), (0,-1)]:
                x, y = i+dx, j+dy
                if 0 <= x < m and 0 <= y < n and not is_visited[x][y]:
                    trap += max(0, height - heightMap[x][y])
                    heappush(heap, [max(height, heightMap[x][y]), x, y])
                    is_visited[x][y] = True

        return trap

[4]:

s = Solution1()

heightMap = [
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]
assert s.trapRainWater(heightMap) == 4

heightMap = [
  [12,13,0,12],
  [13,4,13,12],
  [13,8,10,12],
  [12,13,12,12],
  [13,13,13,13]
]
assert s.trapRainWater(heightMap) == 14

heightMap = [
  [2,2,2,2],
  [2,2,3,4],
  [3,3,3,1],
  [2,3,4,5]
]
assert s.trapRainWater(heightMap) == 0